//! asm code general tools

use crate::data::Segment;
use anyhow::Result;

#[derive(Debug)]
pub(crate) enum AsmArithmeticOpt {
    Add,
    Sub,
    Or,
    And,
}
impl AsmArithmeticOpt {
    pub(crate) fn to_asm(&self) -> &'static str {
        match self {
            Self::Add => "+",
            Self::Sub => "-",
            Self::Or => "|",
            Self::And => "&",
        }
    }
}

/// the asm registor
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum AsmRegistor {
    /// A registor
    A,
    /// D registor
    D,
    /// M registor
    M,
}

impl AsmRegistor {
    /// to asm code
    pub(crate) const fn to_asm(&self) -> &'static str {
        match self {
            Self::A => "A",
            Self::D => "D",
            Self::M => "M",
        }
    }
}

/// the asm registor
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum AsmConditionalJmp {
    /// if comp =0 jump
    Jeq,
    /// if comp >0 jump
    Jgt,
    /// if comp <0 jump
    Jlt,
    /// if comp <> 0 jump
    Jne,
}

impl AsmConditionalJmp {
    /// to asm code
    pub(crate) const fn to_asm(&self) -> &'static str {
        match self {
            Self::Jeq => "JEQ",
            Self::Jgt => "JGT",
            Self::Jlt => "JLT",
            Self::Jne => "JNE",
        }
    }
}

#[derive(Debug)]
pub(crate) struct AsmGenerator {
    vm_file_name: String,
    cache: Vec<String>,
}
impl AsmGenerator {
    pub(crate) const SP: &'static str = "SP";
    /// memory registor
    pub(crate) const MREG: &'static str = AsmRegistor::M.to_asm();
    /// data registor
    pub(crate) const DREG: &'static str = AsmRegistor::D.to_asm();
    /// address registor
    pub(crate) const AREG: &'static str = AsmRegistor::A.to_asm();
    /// create a new generator
    /// # arguments
    /// * `vm_file_name` - the vm file's name
    pub(crate) fn new(vm_file_name: &str) -> Self {
        Self {
            vm_file_name: vm_file_name.into(),
            cache: Vec::new(),
        }
    }
    pub(crate) fn output(&mut self) -> String {
        let code = self.cache.join("\n");
        self.cache.clear();
        code
    }
    /// push the value to the stack.
    ///
    /// # arguments
    /// * `segment` - the segment to read from
    /// * `index` - the index to read from
    pub(crate) fn push(&mut self, segment: &str, index: u32) -> Result<()> {
        let segment = Segment::try_from(segment)?;
        self.push_segment(segment, index);
        Ok(())
    }

    /// push the value to the stack.
    ///
    /// # arguments
    /// * `segment` - the segment to read from
    /// * `index` - the index to read from
    pub(crate) fn push_segment(&mut self, segment: Segment, index: u32) {
        // write to segment direct of constan 0 or 1
        if segment == Segment::Constant && [0, 1].contains(&index) {
            self.add_other(Self::asm_addr_sp());
            self.cache
                .push(Self::asm_assignment(Self::AREG, Self::MREG));
            self.cache
                .push(Self::asm_assignment(Self::MREG, &index.to_string()))
        } else {
            match segment {
                Segment::Constant => self.write_constant_to_dreg(index),
                _ => self.write_non_constant_do_dreg(&segment, index),
            };
            self.write_dreg_to_stack();
        }
        self.sp_inc();
    }

    /// pop value from stack to segment
    ///
    /// # arguments
    /// * `segment` - the segment to write
    /// * `index` - the index to write
    pub(crate) fn pop(&mut self, segment: &str, index: u32) -> Result<()> {
        let segment = Segment::try_from(segment)?;
        self.pop_segment(segment, index);
        Ok(())
    }
    /// load segment address
    fn load_segment_addr(&mut self, segment: Segment, index: u32) {
        // @Segment
        self.add_other(Self::asm_addr(
            &segment.to_asm_signal(&self.vm_file_name, index),
        ));
        if ![Segment::Temp, Segment::Pointer].contains(&segment) {
            // A = M
            self.cache
                .push(Self::asm_assignment(Self::AREG, Self::MREG));
        }
    }
    /// pop value from stack to segment
    ///
    /// # arguments
    /// * `segment` - the segment to write
    /// * `index` - the index to write
    pub(crate) fn pop_segment(&mut self, segment: Segment, index: u32) {
        if segment != Segment::Static {
            if index <= 1 {
                // pop stack value to D registor
                self.pop_to_dreg();
                // @Segment
                self.load_segment_addr(segment, index);
                if index == 1 {
                    // A=A+1
                    self.add_other(Self::asm_assignment(
                        Self::AREG,
                        &Self::asm_arithmetic_add(Self::AREG, "1"),
                    ));
                }
            } else {
                // @Segment
                self.load_segment_addr(segment, index);
                // D=A
                self.cache
                    .push(Self::asm_assignment(Self::DREG, Self::AREG));
                // @index
                self.add_other(Self::asm_addr(&index.to_string()));
                // D=D+A
                self.add_other(Self::asm_assignment(
                    Self::DREG,
                    &Self::asm_arithmetic_add(Self::DREG, Self::AREG),
                ));
                // store the pop value to r14
                // R14=D
                self.write_dreg_to_r14();

                // pop stack value to D registor
                self.pop_to_dreg();

                // @R14
                self.add_other(Self::asm_addr("R14"));
                // A=M
                self.cache
                    .push(Self::asm_assignment(Self::AREG, Self::MREG));
            }
        } else {
            self.pop_to_dreg();
            self.add_other(Self::asm_addr(
                &segment.to_asm_signal(&self.vm_file_name, index),
            ));
        }
        // M=D
        self.add_other(Self::asm_assign_d_to_m());
    }

    //===============

    /// write constant to date register
    ///
    /// # arguments
    /// * `index` - the index to read from
    pub(crate) fn write_constant_to_dreg(&mut self, index: u32) {
        self.cache.extend_from_slice(&[
            Self::asm_addr(index.to_string().as_str()),
            // format!("@{}", index),
            Self::asm_assignment(Self::DREG, Self::AREG),
        ])
    }

    /// write non-constant data to date register
    ///
    /// # arguments
    /// * `segment` - the segment to read from
    /// * `index` - the index to read from
    /// * `vm_file_name` - the vm file's name
    pub(crate) fn write_non_constant_do_dreg(&mut self, segment: &Segment, index: u32) {
        self.add_other(Self::asm_addr(
            &segment.to_asm_signal(&self.vm_file_name, index),
        ));
        /*
        self.add_other(format!(
            "@{}",
            segment.to_asm_signal(&self.vm_file_name, index)
        )); */
        if segment != &Segment::Static {
            if [Segment::Temp, Segment::Pointer].contains(segment) {
                if index > 0 {
                    self.cache
                        .push(Self::asm_assignment(Self::DREG, Self::AREG));
                }
            } else {
                self.add_other(Self::asm_assign_m_to_a_and_d());
            }
            if segment != &Segment::Static && index > 0 {
                self.add_other(Self::asm_addr(index.to_string().as_str()));
                //self.add_other(format!("@{}", index));
                self.add_other(Self::asm_assign_d_add_a_to_a());
            }
        }
        self.add_other(Self::asm_assign_m_to_d());
    }

    /// write date register to sp's value
    pub(crate) fn write_dreg_to_stack(&mut self) {
        self.assign_sp_value_to_areg();
        self.add_other(Self::asm_assign_d_to_m());
    }

    /// load stack pointer  to A register
    pub(crate) fn assign_sp_value_to_areg(&mut self) {
        self.add_other(Self::asm_addr_sp());
        self.add_other(Self::asm_assignment(Self::AREG, Self::MREG));
    }

    /// add other code
    pub(crate) fn add_other(&mut self, code: String) {
        self.cache.push(code);
    }

    /// pop stack's topmost value to d reg
    pub(crate) fn pop_to_dreg(&mut self) {
        self.pop_to_registor(AsmRegistor::D);
    }
    /// pop stack's topmost value to reg
    pub(crate) fn pop_to_registor(&mut self, reg: AsmRegistor) {
        self.papare_pop_from_stack();
        self.add_other(Self::asm_assignment(reg.to_asm(), Self::MREG));
    }
    /// papare to pop from the stack
    ///
    /// # code
    ///
    /// * SP=SP-1
    /// * @SP
    /// * A=M
    ///
    pub(crate) fn papare_pop_from_stack(&mut self) {
        self.sp_dec();
        self.add_other(Self::asm_assignment(Self::AREG, Self::MREG));
    }

    /// push D registor to stack
    pub(crate) fn push_dreg_to_stack(&mut self) {
        self.push_registor_to_stack(AsmRegistor::D);
    }

    /// push registor to stack
    pub(crate) fn push_registor_to_stack(&mut self, reg: AsmRegistor) {
        self.add_other(Self::asm_addr_sp());
        self.add_other(Self::asm_assignment(Self::AREG, Self::MREG));
        self.add_other(Self::asm_assignment(Self::MREG, reg.to_asm()));
        self.sp_inc();
    }
    /// push segment base address to statck
    pub(crate) fn push_segment_base_address_to_stack(&mut self, segment: Segment) {
        self.read_segment_base_address_to_dreg(segment);
        self.push_dreg_to_stack();
    }
    /// read segment base address to d registor
    pub(crate) fn read_segment_base_address_to_dreg(&mut self, segment: Segment) {
        self.add_other(Self::asm_addr(
            &segment.to_asm_signal(&self.vm_file_name, 0),
        ));
        self.add_other(Self::asm_assignment(Self::DREG, Self::MREG));
    }

    /// increment stack pointer
    pub(crate) fn sp_inc(&mut self) {
        self.add_other(Self::asm_addr_sp());
        self.add_other(Self::asm_assignment(
            Self::MREG,
            &Self::asm_arithmetic_add(Self::MREG, "1"),
        ));
        //self.add_other(format!("{0}={0}+1", Self::MREG));
    }

    /// decrement stack pointer
    pub(crate) fn sp_dec(&mut self) {
        self.add_other(Self::asm_addr_sp());
        self.add_other(Self::asm_assignment(
            Self::MREG,
            &Self::asm_arithmetic_sub(Self::MREG, "1"),
        ));
        //self.add_other(format!("{0}={0}-1", Self::MREG));
    }

    /// address relocation and store in A registor and D registor
    ///
    /// # the code
    ///
    /// AD = addr opt value
    ///
    /// for example:
    ///
    /// AD = ARG + 1
    ///
    /// # arguments
    ///
    /// * addr : the base addr
    /// * opt : offset direction +/-
    /// * offset : offset value
    ///
    /// # result
    ///
    /// set A and D registor with the new addess
    ///
    pub(crate) fn addr_offset_and_store_in_a_and_d(
        &mut self,
        addr: &str,
        opt: AsmArithmeticOpt,
        offset: u32,
    ) {
        self.add_other(Self::asm_addr(addr));
        self.add_other(Self::asm_assignment(Self::DREG, Self::MREG));

        if offset > 0 {
            if offset == 1 {
                self.add_other(Self::asm_assignment(
                    &format!("{}{}", Self::AREG, Self::DREG),
                    &Self::asm_arithmetic(Self::DREG, opt, "1"),
                ));
            } else {
                self.add_other(Self::asm_addr(&offset.to_string()));
                self.add_other(Self::asm_assignment(
                    &format!("{}{}", Self::AREG, Self::DREG),
                    &Self::asm_arithmetic(Self::DREG, opt, Self::AREG),
                ));
            }
        }
    }

    /// set d registor to segment base address
    pub(crate) fn set_dreg_to_segment_base_address(&mut self, segment: Segment) {
        self.set_dreg_to_base_address(&segment.to_asm_signal(&self.vm_file_name, 0));
    }
    /// set d registor to segment base address
    pub(crate) fn set_dreg_to_base_address(&mut self, addr: &str) {
        self.add_other(Self::asm_addr(addr));
        self.add_other(Self::asm_assignment(Self::MREG, Self::DREG));
    }
    /// jump to label(function)
    pub(crate) fn jmp_to_label(&mut self, label: &str) {
        self.add_other(Self::asm_addr(label));
        self.add_other(Self::asm_jmp().into());
    }

    pub(crate) fn write_dreg_to_r14(&mut self) {
        self.write_dreg_to_r_reg(14)
    }
    pub(crate) fn write_dreg_to_r15(&mut self) {
        self.write_dreg_to_r_reg(15)
    }

    pub(crate) fn write_dreg_to_r_reg(&mut self, index: u8) {
        self.add_other(Self::asm_addr(&format!("R{index}")));
        // self.add_other(format!("@R{index}"));
        self.add_other(Self::asm_assign_d_to_m());
    }

    /// read r14 value to dreg
    pub(crate) fn read_r15_value_to_dreg(&mut self, offset: u32) {
        self.read_rreg_value_to_dreg(15, offset);
    }
    /// read R registor value to dreg
    pub(crate) fn read_rreg_value_to_dreg(&mut self, index: u8, offset: u32) {
        self.addr_offset_and_store_in_a_and_d(&format!("R{index}"), AsmArithmeticOpt::Sub, offset);
        self.add_other(Self::asm_assign_m_to_d());
    }

    /// set r15 to the segment base addr
    ///
    /// # arguments
    ///
    /// * segment : the updated segment
    /// * offset : segment = r14-offset
    pub(crate) fn set_r15_value_to_segment_base_addr(&mut self, segment: Segment, offset: u32) {
        self.read_r15_value_to_dreg(offset);
        self.set_dreg_to_segment_base_address(segment);
    }

    //==== asm code

    /// the asm code of load addr instruction
    pub(crate) fn asm_addr(addr: &str) -> String {
        format!("@{addr}")
    }

    /// the asm code of  load sp addr instruction
    pub(crate) fn asm_addr_sp() -> String {
        Self::asm_addr(Self::SP)
    }

    // === compose instuction ===

    /// asm code : M=D
    pub(crate) fn asm_assign_d_to_m() -> String {
        Self::asm_assignment(Self::MREG, Self::DREG)
    }
    /// asm code : D=M
    pub(crate) fn asm_assign_m_to_d() -> String {
        Self::asm_assignment(Self::DREG, Self::MREG)
    }

    /// asm code : A=D+A
    pub(crate) fn asm_assign_d_add_a_to_a() -> String {
        Self::asm_assignment(
            Self::AREG,
            &Self::asm_arithmetic_add(Self::DREG, Self::AREG),
        )
    }

    /// asm code : AD=M
    pub(crate) fn asm_assign_m_to_a_and_d() -> String {
        Self::asm_assignment(&format!("{}{}", Self::AREG, Self::DREG), Self::MREG)
    }

    // === core instuction ===

    /// the asm code of assignment instruction
    pub(crate) fn asm_assignment(left: &str, right: &str) -> String {
        format!("{left}={right}")
    }
    /// the asm code of arithmetic add instruction
    pub(crate) fn asm_arithmetic_add(left: &str, right: &str) -> String {
        Self::asm_arithmetic(left, AsmArithmeticOpt::Add, right)
    }
    /// the asm code of arithmetic sub instruction
    pub(crate) fn asm_arithmetic_sub(left: &str, right: &str) -> String {
        Self::asm_arithmetic(left, AsmArithmeticOpt::Sub, right)
    }
    /// the asm code of arithmetic instruction
    pub(crate) fn asm_arithmetic(left: &str, opt: AsmArithmeticOpt, right: &str) -> String {
        format!("{left}{}{right}", opt.to_asm())
    }

    /// the asm code of neg instruction
    ///
    /// the code: -value
    pub(crate) fn asm_arithmetic_neg(value: &str) -> String {
        format!("-{value}")
    }

    /// the asm code of not instruction
    ///
    /// the code: !value
    pub(crate) fn asm_arithmetic_not(value: &str) -> String {
        format!("!{value}")
    }

    /// the asm code of arithmetic and instruction
    /// the code: left&right
    pub(crate) fn asm_arithmetic_and(left: &str, right: &str) -> String {
        Self::asm_arithmetic(left, AsmArithmeticOpt::And, right)
    }

    /// the asm code of arithmetic and instruction
    ///
    /// the code: left|right
    pub(crate) fn asm_arithmetic_or(left: &str, right: &str) -> String {
        Self::asm_arithmetic(left, AsmArithmeticOpt::Or, right)
    }
    /// the asm code of assignment instruction
    ///
    /// the code: 0;JMP
    pub(crate) fn asm_jmp() -> &'static str {
        "0;JMP"
    }

    /// the asm code of assignment instruction
    /// the code: reg;jmp
    pub(crate) fn asm_conditional_jmp(reg: AsmRegistor, jmp: AsmConditionalJmp) -> String {
        format!("{};{}", reg.to_asm(), jmp.to_asm())
    }
    /// the asm code of label
    ///
    /// # the code: (label)
    pub(crate) fn asm_label(label: &str) -> String {
        format!("({label})")
    }

    /// the bootstrap code
    pub(crate) fn asm_bootstrap_code() -> &'static str {
        include_str!("bootstrap.asm")
    }
}

#[cfg(test)]
mod test {
    use crate::data::Segment;

    use super::AsmGenerator;

    #[test]
    fn test_pop_segment_index_2() {
        let mut asm_generator = AsmGenerator::new("");
        asm_generator.pop_segment(Segment::Argument, 2);
        println!("{}", asm_generator.output())
    }

    #[test]
    fn test_pop_segment_index_1() {
        let mut asm_generator = AsmGenerator::new("");
        asm_generator.pop_segment(Segment::Argument, 1);
        println!("{}", asm_generator.output())
    }

    #[test]
    fn test_pop_segment_index_0() {
        let mut asm_generator = AsmGenerator::new("");
        asm_generator.pop_segment(Segment::Argument, 0);
        println!("{}", asm_generator.output())
    }
}
